﻿using System;
using System.Collections.Generic;
using System.Linq;
using System.Text;

using System.Diagnostics.Contracts;
using Microsoft.Z3;

namespace TestGen
{
	public sealed class ConstraintSolve
	{
		/// <summary>
		/// The context in which we build the constraint set.
		/// </summary>
		private readonly Context z3;

		public ConstraintSolve()
		{
			Config p = new Config();
			p.SetParamValue("MODEL", "true");
			this.z3 = new Context(p);
			p.Dispose();
		}

		private Sort MakeIntType()
		{ return this.z3.MkIntSort(); }

		private Term MakeVar(String name, Sort ty)
		{ return this.z3.MkConst(name, ty); }

		private Term MakeIntVar(String name)
		{ return this.MakeVar(name, this.MakeIntType()); }

		private Term MakeIntConst(uint v)
		{ return this.z3.MkNumeral(v, this.MakeIntType()); }

		private Term MakeIntMul(Term t1, Term t2)
		{ return this.z3.MkMul(t1, t2); }

		private Term MakeIntSigmaAdd(Term[] trms)
		{ return this.z3.MkAdd(trms); }

		private Term MakeEquality(Term t1, Term t2)
		{ return this.z3.MkEq(t1, t2); }

		private Term MakeLessEqual(Term t1, Term t2)
		{ return this.z3.MkLe(t1, t2); }

		private Term MakeGreaterEqual(Term t1, Term t2)
		{ return this.z3.MkGe(t1, t2); }

		/// <summary>
		/// Return true if the term <c>f</c> is a redundant constraint wrt. the set we have already added.
		/// </summary>
		private bool IsRedundantConstraint(Term f)
		{
			/* save the current state of the context */
			z3.Push();

			Term not_f = z3.MkNot(f);
			z3.AssertCnstr(not_f);

			Model model = null;
			LBool issat = z3.CheckAndGetModel(out model);

			bool redundant = (issat == LBool.False);

			if(model != null)
				model.Dispose();

			/* restore context */
			z3.Pop(1);

			return redundant;
		}

		/// <summary>
		/// Given a set of constraints generate the implied equalities (equiv classes) and return a mapping for simplifying our constraint system based 
		/// on these equalities.
		/// </summary>
		/// <param name="generatedTermMap">This is a map from the z3 generated terms to the constraint terms</param>
		private Dictionary<TGenVar, TGenAtomTerm> GenerateSimplificationMap(Dictionary<Term, TGenAtomTerm> generatedTermMap)
		{
			Dictionary<TGenVar, TGenAtomTerm> resmap = new Dictionary<TGenVar, TGenAtomTerm>();

			Term[] terms = generatedTermMap.Keys.ToArray();
			uint[] class_ids;

			this.z3.GetImpliedEqualities(terms, out class_ids);
			Dictionary<int, List<TGenAtomTerm>> eqclasses = new Dictionary<int, List<TGenAtomTerm>>();

			for(int i = 0; i < terms.Length; ++i)
			{
				TGenAtomTerm atom = generatedTermMap[terms[i]];
				if(!eqclasses.ContainsKey((int)class_ids[i]))
					eqclasses.Add((int)class_ids[i], new List<TGenAtomTerm>());
				eqclasses[(int)class_ids[i]].Add(atom);
			}

			foreach(List<TGenAtomTerm> eqclass in eqclasses.Values)
			{
				if(eqclass.Count < 2)
				{ ;}
				else
				{
					eqclass.Sort(TGenAtomTerm.CompareAtomTerm);

					for(int j = 1; j < eqclass.Count; ++j)
					{
						if(eqclass[j] is TGenVar)
							resmap.Add((TGenVar)eqclass[j].CopyAtomTerm(), eqclass[0].CopyAtomTerm());
					}
				}
			}

			return resmap;
		}

		/// <summary>
		/// Create a z3 term that represents the given constraint.
		/// </summary>
		public Term MakeTermForTGenConstraint(TGenConstraint tc, Dictionary<Term, TGenAtomTerm> generatedTermMap)
		{
			Contract.Requires((tc.ConstraintOp == ConstraintType.Equality) | (tc.ConstraintOp == ConstraintType.LessEqual) | (tc.ConstraintOp == ConstraintType.GreaterEqual));

			Term tl = this.MakeTermForTGenTerm(tc.LeftTerm, generatedTermMap);
			Term tr = this.MakeTermForTGenTerm(tc.RightTerm, generatedTermMap);

			Term rest = null;
			switch(tc.ConstraintOp)
			{
				case ConstraintType.Equality:
					rest = this.MakeEquality(tl, tr);
					break;
				case ConstraintType.LessEqual:
					rest = this.MakeLessEqual(tl, tr);
					break;
				case ConstraintType.GreaterEqual:
					rest = this.MakeGreaterEqual(tl, tr);
					break;
			}
			return rest;
		}

		/// <summary>
		/// Create a z3 term that represents the given constraint term.
		/// </summary>
		public Term MakeTermForTGenTerm(TGenTerm tt, Dictionary<Term, TGenAtomTerm> generatedTermMap)
		{
			Contract.Requires((tt is TGenConst) | (tt is TGenVar) | (tt is TGenBinMult) | (tt is TGenSum));

			if(tt is TGenConst)
			{
				Term rt = this.MakeIntConst(((TGenConst)tt).ConstVal);
				if(generatedTermMap != null && (!generatedTermMap.ContainsKey(rt)))
					generatedTermMap.Add(rt, (TGenConst)tt);

				return rt;
			}
			else if(tt is TGenVar)
			{
				Term rt = this.MakeIntVar(((TGenVar)tt).VarName);
				if(generatedTermMap != null && (!generatedTermMap.ContainsKey(rt)))
					generatedTermMap.Add(rt, (TGenVar)tt);

				return rt;
			}
			else if(tt is TGenBinMult)
			{
				Term t1 = this.MakeTermForTGenTerm(((TGenBinMult)tt).Val1, generatedTermMap);
				Term t2 = this.MakeTermForTGenTerm(((TGenBinMult)tt).Val2, generatedTermMap);
				return this.MakeIntMul(t1, t2);
			}
			else
			{
				TGenSum tgs = (TGenSum)tt;
				Term[] tarr = new Term[tgs.SumTerms.Length];
				for(int i = 0; i < tgs.SumTerms.Length; ++i)
					tarr[i] = this.MakeTermForTGenTerm(tgs.SumTerms[i], generatedTermMap);

				return this.MakeIntSigmaAdd(tarr);
			}
		}

		/// <summary>
		/// Given a list of constraints we generated for the constructor size info.
		/// Generate a simplifed set to actually emit and for any variables that are now completely constrained emit the value they should be assigned.
		/// </summary>
		public static void GenerateSimplifiedConstraints(List<TGenConstraint> inconstraints, List<TGenConstraint> outconstraints, List<Tuple<String, String>> fixedparams)
		{
			for(int i = 0; i < inconstraints.Count; ++i)
				inconstraints[i] = inconstraints[i].Simplify();

			//generate all the implied equalities
			Dictionary<Term, TGenAtomTerm> generatedTermMap = new Dictionary<Term, TGenAtomTerm>();
			ConstraintSolve cs1 = new ConstraintSolve();
			foreach(TGenConstraint tgc in inconstraints)
			{
				Term gt = cs1.MakeTermForTGenConstraint(tgc, generatedTermMap);
				cs1.z3.AssertCnstr(gt);
			}
			Dictionary<TGenVar, TGenAtomTerm> simplmap = cs1.GenerateSimplificationMap(generatedTermMap);

			//get fully constrained assignments for use in the constructor
			foreach(KeyValuePair<TGenVar, TGenAtomTerm> fasign in simplmap)
				fixedparams.Add(new Tuple<String, String>(fasign.Key.VarName, fasign.Value.GenerateTermOutputForCSharp()));

			//apply the equalities to simplify the constraint system
			List<TGenConstraint> nterms = new List<TGenConstraint>();
			foreach(TGenConstraint tgc in inconstraints)
			{
				TGenTerm leqrw = tgc.LeftTerm.TermRemap(simplmap);
				TGenTerm lsmp = leqrw.Simplify();

				TGenTerm reqrw = tgc.RightTerm.TermRemap(simplmap);
				TGenTerm rsmp = reqrw.Simplify();

				nterms.Add(new TGenConstraint(tgc.ConstraintOp, lsmp, rsmp));
			}

			//take the simplified set of constraints and eliminate all redundant constraints
			ConstraintSolve cs2 = new ConstraintSolve();
			foreach(TGenConstraint tgc in nterms)
			{
				Term constr = cs2.MakeTermForTGenConstraint(tgc, null);
				bool redundant = cs2.IsRedundantConstraint(constr);
				if(!redundant)
				{
					outconstraints.Add(tgc);
					cs2.z3.AssertCnstr(constr);
				}
			}
		}

		//////////////////////////////////////////////////////////////

		public void find_model_example1()
		{
			Console.WriteLine("find_model_example2");

			Term x = this.MakeIntVar("x");
			Term y = this.MakeIntVar("y");
			Term z = this.MakeIntVar("z");
			Term one = this.MakeIntConst(1);
			Term two = this.MakeIntConst(2);

			Term[] terms = new Term[] { x, y, z, one };
			uint[] class_ids;

			Term y_plus_two = this.z3.MkAdd(y, two);
			Term y_plus_one_one = this.z3.MkAdd(this.z3.MkAdd(y, one), one);

			Term ca = this.z3.MkEq(x, y_plus_two);
			Term cb = this.z3.MkEq(z, y_plus_one_one);

			this.z3.AssertCnstr(ca);
			this.z3.AssertCnstr(cb);

			Console.WriteLine("model for: x = y + 2, z = y + 1 + 1");

			this.z3.GetImpliedEqualities(terms, out class_ids);
			for(int i = 0; i < terms.Length; ++i)
			{
				Console.WriteLine("Class {0} |-> {1}", terms[i], class_ids[i]);
			}
		}
	}
}
